# This Python file uses the following encoding: utf-8

import h5py as h5
from template import *




class TriMesh2D:
    def __init__(self):
        pass

    def read_cgns_mesh(self, file_name):
        mesh_file = h5.File(file_name)
        mesh_1 = mesh_file["Mesh_1"]
        smesh_mesh = mesh_1["SMESH_Mesh"]

        edge_element = list(smesh_mesh.keys())[3]
        tri_element  = list(smesh_mesh.keys())[4]

        self.vertex_x = smesh_mesh["GridCoordinates"]["CoordinateX"][" data"][:]
        self.vertex_y = smesh_mesh["GridCoordinates"]["CoordinateY"][" data"][:]
        self.vertex_z = smesh_mesh["GridCoordinates"]["CoordinateZ"][" data"][:]

        self.bar_elements = smesh_mesh[edge_element]["ElementConnectivity"][" data"][:]
        self.bar_elements = self.bar_elements.reshape((int(self.bar_elements.shape[0]/2), 2)) - 1

        self.tri_elements = smesh_mesh[tri_element]["ElementConnectivity"][" data"][:]
        self.tri_elements = self.tri_elements.reshape((int(self.tri_elements.shape[0]/3), 3)) - 1


    def plot_mesh(self, index0 = 0, index1 = 0):

        fig=plt.figure(figsize=(10,8))
        fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.4)

        #------ potential ---------------
        ax0 = fig.add_subplot(1,1,1)

        ax0.triplot(self.vertex_x, self.vertex_y, self.tri_elements, lw=1.0)
        #ax0.triplot(self.vertex_x, self.vertex_y, self.tri_elements[0:1, :], lw=1.0)

        if index1 > index0:
            edge_points = np.zeros((index1 - index0, 2))
            for i in range(index0, index1):
                edge_points[i-index0, 0] = self.vertex_x[self.bar_elements[i, 0]]
                edge_points[i-index0, 1] = self.vertex_y[self.bar_elements[i, 0]]
            ax0.scatter(edge_points[:, 0], edge_points[:, 1])
        ax0.set_aspect('equal')


        fig.savefig("mesh.svg", dpi = 300)
        plt.show()





'''
if__name__ == "__main__":
    mesh = TriMesh2D ()
    mesh.read_cgns_mesh("mesh100.cgns")
    mesh.plot_mesh()
'''
